# -*- coding: utf-8 -*-
"""
-------------------------------------------------
   File Name：     visual_utils
   Description :
   Author :       Flyoung
   date：          2023/3/20
-------------------------------------------------
   Change Activity:
                   2023/3/20:
-------------------------------------------------
"""
import json
from typing import Dict, AnyStr


def state_transition_table_visual(d_tran: Dict, start_state=None, end_states=None, title: AnyStr = None):
    import networkx as nx
    import matplotlib.pyplot as plt
    # 创建有向图
    G = nx.DiGraph()

    def get_node_label(node_):
        if node_ != start_state and node_ not in end_states:
            return node_
        elif node_ == start_state:
            return f'Start ({node_})'
        elif node_ in end_states:
            return f'End ({node_})'

    # 遍历所有节点和边来添加到有向图中
    edges = [
        (node, target_node, {'label': edge_label}) for node, edges in d_tran.items()
        for edge_label, target_nodes in edges.items() for target_node in target_nodes
    ]
    # 将重边的标签合并
    added_edge = dict()
    for edge in edges:
        from_to = edge[:2]
        if from_to in added_edge:
            e = added_edge[from_to]
            e[-1]['label'] += f' | {edge[-1]["label"]}'
        else:
            added_edge[from_to] = edge
    G.add_edges_from(added_edge.values())
    # 添加节点
    for node in d_tran.keys():
        if start_state and end_states:
            G.add_node(node, label=get_node_label(node))
        else:
            G.add_node(node, label=node)

    # 绘制有向图
    plt.figure(figsize=(10, 10))
    pos = nx.circular_layout(G)

    # 定义标签位置
    label_pos = {}
    for node, coords in pos.items():
        # 在节点左侧略微偏移
        label_pos[node] = (coords[0] - 0.1, coords[1])

    # 绘制节点
    nx.draw_networkx_nodes(G, pos, node_color='skyblue', alpha=0.5)
    # 绘制边
    nx.draw_networkx_edges(G, pos)
    # 绘制节点标签
    nodes_label = {n: G.nodes[n]['label'] for n in G.nodes()}
    nx.draw_networkx_labels(
        G, label_pos,
        labels=nodes_label,
    )
    # 绘制边的标签
    nx.draw_networkx_edge_labels(
        G, pos, edge_labels=nx.get_edge_attributes(G, 'label'),
        font_color='red',
        font_weight=2,
        font_size=15
    )
    if title:
        plt.title(title, fontsize=20, fontweight='bold')
    plt.tight_layout()
    plt.axis('off')
    plt.show()


def show_dict(d: Dict, indent: int = 4, split_str_len: int = 50, title: AnyStr = None):
    split_str = '=' * split_str_len
    if title:
        print(f'{split_str} {title} {split_str}')
    print(json.dumps(d, indent=indent))
    print(split_str * 2)
